import geopandas as gpd
import numpy as np
import json
from pykrige.ok import OrdinaryKriging
import pandas as pd
import rasterio
from rasterio.mask import mask
from rasterio.transform import from_bounds
from shapely.geometry import mapping

import plotnine
from plotnine import *
from shapely import Point

from test4 import load_config


def crop_with_geojson(data, region, transform):
    import numpy as np
    with rasterio.io.MemoryFile() as memfile:
        with memfile.open(
            driver='GTiff',
            height=data.shape[0],
            width=data.shape[1],
            count=1,
            crs='EPSG:4326',
            dtype='float32',
            transform=transform
        ) as dataset:
            dataset.write(data, 1)
            out_image, _ = mask(dataset, [mapping(region)], crop=False, all_touched=True)
            out_image[out_image == 0] = np.nan
            return out_image[0]
def create_grid(bounding_box, grid_size=1024):
    """生成插值网格"""
    grid_lon = np.linspace(bounding_box["bottom_left"][0], bounding_box["top_right"][0], grid_size)
    grid_lat = np.linspace(bounding_box["bottom_left"][1], bounding_box["top_right"][1], grid_size)
    return np.meshgrid(grid_lon, grid_lat)
# 读取 GeoJSON 文件
gdf = gpd.read_file('100000.geo.json')
gdf = gdf.to_crs(epsg=4326)
gdf_bounds = gdf.total_bounds
region = gdf.geometry.unary_union
#
# grid_lon = np.linspace(gdf_bounds[0], gdf_bounds[2], 400)
# grid_lat = np.linspace(gdf_bounds[1], gdf_bounds[3], 400)
CONFIG = load_config("application.yaml")

bounding_box = {
    "bottom_left": (CONFIG['bounding_box'][0], CONFIG['bounding_box'][1]),  # 左下角坐标 , 3.513421, 73.388672
    "top_right": (CONFIG['bounding_box'][2], CONFIG['bounding_box'][3]),  # 右上角坐标  53.696706, 136.142578
}
# grid_lon = np.linspace(china_bounds[0], china_bounds[2], 400)
# grid_lat = np.linspace(china_bounds[1], china_bounds[3], 400)
grid_x, grid_y = create_grid(bounding_box, 400)

grid_lon = grid_x[0]
grid_lat = grid_y[:, 0]
with open('data2.json', 'r') as f:
    data_list = json.load(f)
data = np.array(data_list)
points_lat = data[:, 0]
points_lon = data[:, 1]
points_aqi = data[:, 2]

OK = OrdinaryKriging(points_lon, points_lat, points_aqi, variogram_model='gaussian', nlags=6)
z1, ss1 = OK.execute('grid', grid_lon, grid_lat)

# 计算 transform
transform = from_bounds(
    gdf_bounds[0], gdf_bounds[1], gdf_bounds[2], gdf_bounds[3],
    z1.shape[1], z1.shape[0]
)

# GeoJSON 区域裁剪
z1_crop = crop_with_geojson(z1, region, transform)
# 翻转
z1_crop = np.flipud(z1_crop)
# 转换成网格
xgrid, ygrid = np.meshgrid(grid_lon, grid_lat)
df_grid = pd.DataFrame({
    'long': xgrid.flatten(),
    'lat': ygrid.flatten(),
    'Krig_gaussian': z1_crop.flatten()
})

# 绘图
Krig_inter_grid = (ggplot() +
                   geom_tile(df_grid, aes(x='long', y='lat', fill='Krig_gaussian'), size=0.1) +
                   geom_map(gdf, fill='none', color='none', size=1) +
                   scale_fill_gradientn(colors=['#9cd84e', '#facf39', '#f99049',
                                                '#f65e5f', '#a070b6', '#a06a7b'],
                                        breaks=[50, 100, 150, 200, 300, 400],
                                        name='AQI') +
                   theme(
                       text=element_text(family="Consolas", size=12),
                       legend_position='none',
                       axis_title=element_blank(),
                       axis_text=element_blank(),
                       axis_ticks=element_blank(),
                       panel_background=element_blank(),
                       axis_ticks_major_x=element_blank(),
                       axis_ticks_major_y=element_blank(),
                       panel_grid_major=element_blank(),
                       panel_grid_minor=element_blank()
                   ))

Krig_inter_grid.save("/home/hui/Documents/warehouse/air_monitor/air-monitor/air-monitor-ui/src/assets/images/img-2025-01-10-16.png", dpi=300, bbox_inches='tight',  # 紧凑布局
                     transparent=True,  # 透明背景
                     pad_inches=0)